package tools

import (
	"encoding/json"
	"fmt"
	"regexp"
	"time"

	"workflow-gui/internal/ai"
	"workflow-gui/internal/database"
	"workflow-gui/internal/protocol"
	"workflow-gui/internal/workflow"
)

// AIIterateHandler handles AI iteration requests
type AIIterateHandler struct {
	db           *database.DB
	aiClient     *ai.AIClient
	promptEngine *ai.PromptEngine
}

// NewAIIterateHandler creates a new AI iterate handler
func NewAIIterateHandler(db *database.DB, aiClient *ai.AIClient, promptEngine *ai.PromptEngine) *AIIterateHandler {
	return &AIIterateHandler{
		db:           db,
		aiClient:     aiClient,
		promptEngine: promptEngine,
	}
}

// Adjustment represents a single field adjustment
type Adjustment struct {
	Field     string      `json:"field"`
	OldValue  interface{} `json:"old_value"`
	NewValue  interface{} `json:"new_value"`
	Reasoning string      `json:"reasoning"`
}

// IterateResponse represents the AI iterate response
type IterateResponse struct {
	Adjustments      []Adjustment `json:"adjustments"`
	CascadingEffects []string     `json:"cascading_effects,omitempty"`
}

// HandleAIIterateStage1 handles Stage 1 iteration request
func (h *AIIterateHandler) HandleAIIterateStage1(args map[string]interface{}) (protocol.ToolResult, error) {
	// Extract arguments
	changeID, ok := args["change_id"].(string)
	if !ok || changeID == "" {
		return protocol.NewErrorResult("Missing required argument: change_id"), nil
	}

	userFeedback, ok := args["user_feedback"].(string)
	if !ok || userFeedback == "" {
		return protocol.NewErrorResult("Missing required argument: user_feedback"), nil
	}

	// Load project from database
	dao := database.NewProjectDAO(h.db)
	project, err := dao.GetByChangeID(changeID)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to load project: %v", err)), nil
	}

	// Check workflow_stage
	if project.Meta.WorkflowStage != "ai_reviewing_stage1" && project.Meta.WorkflowStage != "stage1_confirmed" {
		return protocol.NewErrorResult(fmt.Sprintf("Invalid workflow_stage: expected 'ai_reviewing_stage1' or 'stage1_confirmed', got '%s'", project.Meta.WorkflowStage)), nil
	}

	// Build iterate prompt
	prompt, err := h.buildIteratePrompt("stage1", project, userFeedback)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to build prompt: %v", err)), nil
	}

	// Call AI
	response, err := h.aiClient.CallClaude(prompt)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("AI call failed: %v", err)), nil
	}

	// Parse response
	iterateResp, err := h.parseIterateResponse(response)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to parse AI response: %v", err)), nil
	}

	// Validate adjustments
	if err := ai.ValidateIterationOutput(iterateResp.Adjustments); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Validation failed: %v", err)), nil
	}

	// Apply adjustments
	if err := h.applyAdjustments(changeID, "stage1", iterateResp.Adjustments); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to apply adjustments: %v", err)), nil
	}

	// Save iteration record
	if err := h.saveIterationRecord(changeID, "stage1", userFeedback, iterateResp); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to save iteration record: %v", err)), nil
	}

	return protocol.NewSuccessResult(fmt.Sprintf("✅ Stage 1 迭代成功！\n\n已应用 %d 个调整\n\n状态保持为: %s", len(iterateResp.Adjustments), project.Meta.WorkflowStage)), nil
}

// HandleAIIterateStage2 handles Stage 2 iteration request
func (h *AIIterateHandler) HandleAIIterateStage2(args map[string]interface{}) (protocol.ToolResult, error) {
	// Extract arguments
	changeID, ok := args["change_id"].(string)
	if !ok || changeID == "" {
		return protocol.NewErrorResult("Missing required argument: change_id"), nil
	}

	userFeedback, ok := args["user_feedback"].(string)
	if !ok || userFeedback == "" {
		return protocol.NewErrorResult("Missing required argument: user_feedback"), nil
	}

	// Load project from database
	dao := database.NewProjectDAO(h.db)
	project, err := dao.GetByChangeID(changeID)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to load project: %v", err)), nil
	}

	// Check workflow_stage
	if project.Meta.WorkflowStage != "ai_reviewing_stage2" && project.Meta.WorkflowStage != "stage2_confirmed" {
		return protocol.NewErrorResult(fmt.Sprintf("Invalid workflow_stage: expected 'ai_reviewing_stage2' or 'stage2_confirmed', got '%s'", project.Meta.WorkflowStage)), nil
	}

	// Build iterate prompt
	prompt, err := h.buildIteratePrompt("stage2", project, userFeedback)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to build prompt: %v", err)), nil
	}

	// Call AI
	response, err := h.aiClient.CallClaude(prompt)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("AI call failed: %v", err)), nil
	}

	// Parse response
	iterateResp, err := h.parseIterateResponse(response)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to parse AI response: %v", err)), nil
	}

	// Validate adjustments
	if err := ai.ValidateIterationOutput(iterateResp.Adjustments); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Validation failed: %v", err)), nil
	}

	// Apply adjustments
	if err := h.applyAdjustments(changeID, "stage2", iterateResp.Adjustments); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to apply adjustments: %v", err)), nil
	}

	// Save iteration record
	if err := h.saveIterationRecord(changeID, "stage2", userFeedback, iterateResp); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to save iteration record: %v", err)), nil
	}

	return protocol.NewSuccessResult(fmt.Sprintf("✅ Stage 2 迭代成功！\n\n已应用 %d 个调整\n\n状态保持为: %s", len(iterateResp.Adjustments), project.Meta.WorkflowStage)), nil
}

// HandleAIIterateStage3 handles Stage 3 iteration request
func (h *AIIterateHandler) HandleAIIterateStage3(args map[string]interface{}) (protocol.ToolResult, error) {
	// Extract arguments
	changeID, ok := args["change_id"].(string)
	if !ok || changeID == "" {
		return protocol.NewErrorResult("Missing required argument: change_id"), nil
	}

	userFeedback, ok := args["user_feedback"].(string)
	if !ok || userFeedback == "" {
		return protocol.NewErrorResult("Missing required argument: user_feedback"), nil
	}

	// Load project from database
	dao := database.NewProjectDAO(h.db)
	project, err := dao.GetByChangeID(changeID)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to load project: %v", err)), nil
	}

	// Check workflow_stage
	if project.Meta.WorkflowStage != "completed" {
		return protocol.NewErrorResult(fmt.Sprintf("Invalid workflow_stage: expected 'completed', got '%s'", project.Meta.WorkflowStage)), nil
	}

	// Build iterate prompt
	prompt, err := h.buildIteratePrompt("stage3", project, userFeedback)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to build prompt: %v", err)), nil
	}

	// Call AI
	response, err := h.aiClient.CallClaude(prompt)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("AI call failed: %v", err)), nil
	}

	// Parse response
	iterateResp, err := h.parseIterateResponse(response)
	if err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to parse AI response: %v", err)), nil
	}

	// Validate adjustments
	if err := ai.ValidateIterationOutput(iterateResp.Adjustments); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Validation failed: %v", err)), nil
	}

	// Apply adjustments
	if err := h.applyAdjustments(changeID, "stage3", iterateResp.Adjustments); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to apply adjustments: %v", err)), nil
	}

	// Save iteration record
	if err := h.saveIterationRecord(changeID, "stage3", userFeedback, iterateResp); err != nil {
		return protocol.NewErrorResult(fmt.Sprintf("Failed to save iteration record: %v", err)), nil
	}

	return protocol.NewSuccessResult(fmt.Sprintf("✅ Stage 3 迭代成功！\n\n已应用 %d 个调整\n\n状态保持为: completed", len(iterateResp.Adjustments))), nil
}

// buildIteratePrompt builds the iterate prompt
func (h *AIIterateHandler) buildIteratePrompt(stage string, project *workflow.WorkflowData, userFeedback string) (string, error) {
	var templateName string
	var data interface{}

	switch stage {
	case "stage1":
		templateName = "stage1_iterate"
		data = ai.PromptData{
			Summary:    project.Proposal.Summary,
			Motivation: project.Proposal.Motivation,
			ChangeType: project.Proposal.Scope.ChangeType,
		}
	case "stage2":
		templateName = "stage2_iterate"
		data = ai.Stage2PromptData{
			Summary:              project.Proposal.Summary,
			Motivation:           project.Proposal.Motivation,
			Priority:             project.Proposal.Scope.Priority,
			TargetUsers:          project.Proposal.Scope.TargetUsers,
			KeyFeatures:          project.Proposal.Scope.KeyFeatures,
			AffectedCapabilities: project.Proposal.Scope.AffectedCapabilities,
			BreakingChanges:      project.Proposal.Scope.BreakingChanges,
		}
	case "stage3":
		templateName = "stage3_iterate"
		// Convert workflow.Requirement to ai.Requirement
		aiRequirements := make([]ai.Requirement, len(project.Spec.Requirements))
		for i, req := range project.Spec.Requirements {
			aiRequirements[i] = ai.Requirement{
				ReqID:       req.ID,
				Name:        req.Name,
				Description: req.Description,
			}
		}
		data = ai.Stage3PromptData{
			Requirements: aiRequirements,
		}
	default:
		return "", fmt.Errorf("invalid stage: %s", stage)
	}

	// Render prompt
	prompt, err := h.promptEngine.RenderPrompt(templateName, data)
	if err != nil {
		return "", fmt.Errorf("failed to render prompt: %w", err)
	}

	// Append user feedback
	prompt = fmt.Sprintf("%s\n\n---\n\n## 用户反馈\n\n%s\n\n请根据用户反馈调整相应字段，并说明调整理由。", prompt, userFeedback)

	return prompt, nil
}

// parseIterateResponse parses AI iterate response
func (h *AIIterateHandler) parseIterateResponse(response string) (*IterateResponse, error) {
	// Extract JSON from markdown code blocks
	re := regexp.MustCompile(`(?s)` + "```json\\s*(.+?)\\s*```")
	matches := re.FindStringSubmatch(response)

	var jsonStr string
	if len(matches) > 1 {
		jsonStr = matches[1]
	} else {
		jsonStr = response
	}

	// Parse JSON
	var result IterateResponse
	if err := json.Unmarshal([]byte(jsonStr), &result); err != nil {
		return nil, fmt.Errorf("failed to parse JSON: %w", err)
	}

	return &result, nil
}

// applyAdjustments applies adjustments to database
func (h *AIIterateHandler) applyAdjustments(changeID string, stage string, adjustments []Adjustment) error {
	// Begin transaction
	tx, err := h.db.Begin()
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	// Get project ID
	var projectID int
	err = tx.QueryRow("SELECT id FROM projects WHERE change_id = ?", changeID).Scan(&projectID)
	if err != nil {
		return fmt.Errorf("failed to get project ID: %w", err)
	}

	// Apply adjustments based on stage
	for _, adj := range adjustments {
		switch stage {
		case "stage1":
			if err := h.applyStage1Adjustment(tx, projectID, adj); err != nil {
				return fmt.Errorf("failed to apply stage1 adjustment: %w", err)
			}
		case "stage2":
			if err := h.applyStage2Adjustment(tx, projectID, adj); err != nil {
				return fmt.Errorf("failed to apply stage2 adjustment: %w", err)
			}
		case "stage3":
			if err := h.applyStage3Adjustment(tx, projectID, adj); err != nil {
				return fmt.Errorf("failed to apply stage3 adjustment: %w", err)
			}
		}
	}

	// Commit transaction
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}

	return nil
}

// applyStage1Adjustment applies Stage 1 adjustment
func (h *AIIterateHandler) applyStage1Adjustment(tx interface{}, projectID int, adj Adjustment) error {
	// Type assert tx to have Exec method
	type Executor interface {
		Exec(query string, args ...interface{}) (interface{}, error)
	}
	executor, ok := tx.(Executor)
	if !ok {
		return fmt.Errorf("invalid transaction type")
	}

	// Update proposals table based on field
	switch adj.Field {
	case "priority":
		_, err := executor.Exec("UPDATE proposals SET priority = ? WHERE project_id = ?", adj.NewValue, projectID)
		return err
	case "target_users", "breaking_changes":
		// Update scope_json
		_, err := executor.Exec("UPDATE proposals SET scope_json = json_set(scope_json, '$."+adj.Field+"', ?) WHERE project_id = ?", adj.NewValue, projectID)
		return err
	default:
		return fmt.Errorf("unknown field: %s", adj.Field)
	}
}

// applyStage2Adjustment applies Stage 2 adjustment
func (h *AIIterateHandler) applyStage2Adjustment(tx interface{}, projectID int, adj Adjustment) error {
	// Type assert tx to have Exec method
	type Executor interface {
		Exec(query string, args ...interface{}) (interface{}, error)
	}
	executor, ok := tx.(Executor)
	if !ok {
		return fmt.Errorf("invalid transaction type")
	}

	// Update requirements table based on field
	switch adj.Field {
	case "name", "description", "operation":
		// Assuming we're updating a specific requirement by ID
		// In real implementation, we'd need to identify which requirement to update
		_, err := executor.Exec("UPDATE requirements SET "+adj.Field+" = ? WHERE project_id = ? LIMIT 1", adj.NewValue, projectID)
		return err
	default:
		return fmt.Errorf("unknown field: %s", adj.Field)
	}
}

// applyStage3Adjustment applies Stage 3 adjustment
func (h *AIIterateHandler) applyStage3Adjustment(tx interface{}, projectID int, adj Adjustment) error {
	// Type assert tx to have Exec method
	type Executor interface {
		Exec(query string, args ...interface{}) (interface{}, error)
	}
	executor, ok := tx.(Executor)
	if !ok {
		return fmt.Errorf("invalid transaction type")
	}

	// Update scenarios table based on field
	switch adj.Field {
	case "name", "when_condition", "then_result":
		// Assuming we're updating a specific scenario by ID
		// In real implementation, we'd need to identify which scenario to update
		_, err := executor.Exec("UPDATE scenarios SET "+adj.Field+" = ? WHERE requirement_id IN (SELECT id FROM requirements WHERE project_id = ?) LIMIT 1", adj.NewValue, projectID)
		return err
	default:
		return fmt.Errorf("unknown field: %s", adj.Field)
	}
}

// saveIterationRecord saves iteration record to database
func (h *AIIterateHandler) saveIterationRecord(changeID string, stage string, feedback string, iterateResp *IterateResponse) error {
	// Begin transaction
	tx, err := h.db.Begin()
	if err != nil {
		return fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer tx.Rollback()

	// Get project ID
	var projectID int
	err = tx.QueryRow("SELECT id FROM projects WHERE change_id = ?", changeID).Scan(&projectID)
	if err != nil {
		return fmt.Errorf("failed to get project ID: %w", err)
	}

	// Marshal adjustments to JSON
	adjustmentsJSON, err := json.Marshal(iterateResp.Adjustments)
	if err != nil {
		return fmt.Errorf("failed to marshal adjustments: %w", err)
	}

	// Get current iteration number
	var iterationNumber int
	err = tx.QueryRow("SELECT COALESCE(MAX(iteration_number), 0) FROM requirement_iterations WHERE project_id = ?", projectID).Scan(&iterationNumber)
	if err != nil {
		return fmt.Errorf("failed to get iteration number: %w", err)
	}
	iterationNumber++

	// Insert iteration record
	_, err = tx.Exec(`
		INSERT INTO requirement_iterations (project_id, iteration_number, stage, user_feedback, ai_adjustments_json, created_at)
		VALUES (?, ?, ?, ?, ?, ?)`,
		projectID, iterationNumber, stage, feedback, string(adjustmentsJSON), time.Now())
	if err != nil {
		return fmt.Errorf("failed to insert iteration record: %w", err)
	}

	// Commit transaction
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit transaction: %w", err)
	}

	return nil
}
